package tcp

import (
	"encoding/binary"
	"fmt"
	"github.com/kataras/golog"
	"go.uber.org/ratelimit"
	"golang.org/x/net/context"
	"math/rand"
	"net"
	"nprobe/core"
	"nprobe/enviroments"
	"nprobe/utils/collections"
	"nprobe/utils/routines"
	"time"
)

type TCPSender struct {
	core.BaseSender
}

func NewTCPSender(
	ctx context.Context, name string, id int, dstIPs []string, task *core.ProbeTask, salts map[int][]byte, receiverCnt int,
) *TCPSender {

	ppRBuf := core.NewPPRBuf(task.SpanMs, task.TimeoutMs)
	copiedSalts := collections.DeepCopyIntMap(salts)

	return &TCPSender{
		BaseSender: core.BaseSender{
			Ctx:         ctx,
			Name:        name,
			ID:          id,
			DstIPs:      dstIPs,
			Task:        task,
			RBuf:        ppRBuf,
			Salts:       copiedSalts,
			ReceiverCnt: receiverCnt,
		},
	}
}

func (s *TCPSender) StartSender() error {
	routines.GoWithRecover(s.Name, s.doSend)
	return nil
}

func (s *TCPSender) doSend() {

	golog.Debugf("run sender. name: %s", s.Name)

	rc, err := core.OpenSenderIP4RawPacketConn("ip4:tcp", "0.0.0.0")
	if err != nil {
		return
	}
	s.Conn = rc

	limiter := ratelimit.New(s.Task.PPS, ratelimit.Per(time.Second))
	var cnt uint64
	var seq uint16
	s.CurrSpanRecord = core.NewSpanRecord(0)
	for {
		for _, srcIP := range s.Task.SrcIPs {
			for dstPort := s.Task.DstPortMin; dstPort <= s.Task.DstPortMax; dstPort++ {
				for srcPort := s.Task.SrcPortMin; srcPort <= s.Task.SrcPortMax; srcPort++ {
					limiter.Take()

					seq++
					ipID := uint16(rand.Intn(1 << 16))
					for _, dstIP := range s.DstIPs {
						select {
						case <-s.Ctx.Done():
							s.StopSender()
							return
						default:
							cnt++
							s.sendOnePkt(srcIP, dstIP, srcPort, dstPort, ipID, seq, cnt)
						}
					}
				}
			}

		}
	}
}

func (s *TCPSender) sendOnePkt(srcIP string, dstIP string, srcPort uint16, dstPort uint16, ipID uint16, seq uint16, cnt uint64) error {
	// payload
	pid := enviroments.PID
	receiverID := cnt % uint64(s.ReceiverCnt)
	txTs := time.Now().UnixNano()
	payload := s.Salts[int(seq)%4]
	binary.BigEndian.PutUint16(payload[8:], uint16(receiverID))
	binary.BigEndian.PutUint16(payload[10:], uint16(pid))
	binary.BigEndian.PutUint64(payload[12:], uint64(txTs))
	binary.BigEndian.PutUint16(payload[20:], seq)

	tcpSeq := uint32(txTs)
	msg, err := core.EncodeIPTCPPacket(srcIP, dstIP, ipID, s.Task.TOS, s.Task.TTL, srcPort, dstPort, tcpSeq, payload)
	if err != nil {
		return err
	}

	err = s.Conn.SetWriteDeadline(time.Now().Add(500 * time.Millisecond))
	if err != nil {
		return err
	}
	_, err = s.Conn.WriteToIP(msg, &net.IPAddr{IP: net.ParseIP(dstIP)})
	if err != nil {
		return err
	}

	s.recordPkt(srcIP, dstIP, srcPort, dstPort, seq, txTs)
	return nil
}

func (s *TCPSender) recordPkt(srcIP string, dstIP string, srcPort uint16, dstPort uint16, seq uint16, txTs int64) {
	spanID := txTs / (s.Task.SpanMs * int64(time.Millisecond))
	golog.Debugf("%s: %s:%d->%s:%d, seq: %d, spanID: %d, txTs: %d", s.Name, srcIP, srcPort, dstIP, dstPort, seq, spanID, txTs)

	if spanID > s.CurrSpanRecord.SpanID {
		s.CurrSpanRecord = core.NewSpanRecord(spanID)
		s.RBuf.Enqueue(s.CurrSpanRecord)
	}

	key := fmt.Sprintf("%s-%d", dstIP, seq)
	s.CurrSpanRecord.PktRecordMap[key] = &core.PktRecord{
		Key:     key,
		SrcIP:   srcIP,
		DstIP:   dstIP,
		SrcPort: srcPort,
		DstPort: dstPort,
		Proto:   s.Task.Proto,
		Tos:     s.Task.TOS,
		TxTs:    txTs,
	}
}
